package site.higgs.limiter.ratelimiter.support.guava;

import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import org.redisson.api.RateIntervalUnit;
import site.higgs.limiter.ratelimiter.RateLimiter;

import java.util.Objects;
import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;

/**
 * Created by caocg on 2018/9/23.
 */
public class GuavaRateLimiter extends RateLimiter {


    private String name;

    private LoadingCache<Object, com.google.common.util.concurrent.RateLimiter> caches;

    private Timer timer;


    public GuavaRateLimiter(String name, Config config) {
        this(name, config.getSize(), config.getDuration(), config.getTimeUnit(), config.getTimerduration());
    }

    public GuavaRateLimiter(String name) {
        this(name, 2 << 20, 300, TimeUnit.SECONDS, 1800000);
    }

    public GuavaRateLimiter(String name, int size, long duration, TimeUnit timeUnit, long timerduration) {

        this.caches = CacheBuilder.newBuilder()
                .expireAfterAccess(duration, timeUnit)
                .maximumSize(size).build(new CacheLoader<Object, com.google.common.util.concurrent.RateLimiter>() {
                    @Override
                    public com.google.common.util.concurrent.RateLimiter load(Object key) throws Exception {
                        KeyWrapper keyWrapper = (KeyWrapper) key;
                        double permitsPerSecond;
                        if (keyWrapper.ppsTimeUnit == RateIntervalUnit.MINUTES) {
                            permitsPerSecond = keyWrapper.pps / (keyWrapper.intervalTime * 60);
                        } else if (keyWrapper.ppsTimeUnit == RateIntervalUnit.HOURS) {
                            permitsPerSecond = keyWrapper.pps / (keyWrapper.intervalTime * 60 * 60);
                        } else if (keyWrapper.ppsTimeUnit == RateIntervalUnit.DAYS) {
                            permitsPerSecond = keyWrapper.pps / (keyWrapper.intervalTime * 24 * 60 * 60);
                        } else {
                            permitsPerSecond = keyWrapper.pps / keyWrapper.intervalTime;
                        }
                        return com.google.common.util.concurrent.RateLimiter.create(permitsPerSecond);

                    }
                });
        this.name = name;
        this.timer = new Timer();
        this.timer.schedule(new TimerTask() {
            @Override
            public void run() {
                caches.cleanUp();

            }
        }, timerduration, timerduration);

    }


    @Override
    public boolean acquire(Object key, double pps,long intervalTime, RateIntervalUnit ppsTimeUnit) {
        KeyWrapper keyWrapper = new KeyWrapper(key, pps,intervalTime,ppsTimeUnit);
        try {
            com.google.common.util.concurrent.RateLimiter rateLimiter = this.caches.get(keyWrapper);
            return rateLimiter.tryAcquire();
        } catch (ExecutionException e) {
            return false;
        }
    }

    @Override
    public String getName() {
        return this.name;
    }

    public static class KeyWrapper {

        private Object key;

        private double pps;

        long intervalTime;

        private RateIntervalUnit ppsTimeUnit;

        public KeyWrapper(Object key, double pps, long intervalTime, RateIntervalUnit ppsTimeUnit) {
            this.key = key;
            this.pps = pps;
            this.intervalTime = intervalTime;
            this.ppsTimeUnit = ppsTimeUnit;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            KeyWrapper that = (KeyWrapper) o;
            return Double.compare(that.pps, pps) == 0 && intervalTime == that.intervalTime && key.equals(that.key) && ppsTimeUnit == that.ppsTimeUnit;
        }

        @Override
        public int hashCode() {
            return Objects.hash(key, pps, intervalTime, ppsTimeUnit);
        }
    }
}
